# 在神经网络中的应用
# 通过配置文件 来构建网络模型


import torch.nn as nn


class LayerRegistry:  # 实现一个注册器
    def __init__(self):
        self.layers = dict()

    def register(self, layer_name):
        # 让装饰器接受 layer 参数
        def decorator(layer):
            # 开始注册
            self.layers[layer_name] = layer
            return layer  # 返回注册的层
        return decorator

    def get_layer(self, layer_name):
        if layer_name in self.layers:
            return self.layers[layer_name]
        else:
            raise KeyError(f"未注册的层 '{layer_name}'.")


# 实例化自定义层注册器
layer_register = LayerRegistry()


@layer_register.register("ConvBNReLU")
class ConvBNReLU(nn.Module):  # 自定义层类
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ConvBNReLU, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

    def forward(self, x):
        return self.layers(x)


# 继续注册其他模块
@layer_register.register("BatchNorm2d")
class BatchNorm2d(nn.Module):
    def __init__(self, num_features, *args, **kwargs):
        super(BatchNorm2d, self).__init__()
        self.bn = nn.BatchNorm2d(num_features, *args, **kwargs)

    def forward(self, x):
        return self.bn(x)


@layer_register.register("ReLU")
class ReLU(nn.Module):
    def __init__(self, *args, **kwargs):
        super(ReLU, self).__init__()
        self.relu = nn.ReLU(*args, **kwargs)

    def forward(self, x):
        return self.relu(x)


@layer_register.register("MaxPooling")
class MaxPooling(nn.Module):
    def __init__(self, kernel_size, stride=1, padding=0):
        super(MaxPooling, self).__init__()
        self.maxpool = nn.MaxPool2d(kernel_size, stride=stride, padding=padding)

    def forward(self, x):
        return self.maxpool(x)


@layer_register.register("AvgPooling")
class AvgPooling(nn.Module):
    def __init__(self, kernel_size, stride=1, padding=0):
        super(AvgPooling, self).__init__()
        self.avgpool = nn.AvgPool2d(kernel_size, stride=stride, padding=padding)

    def forward(self, x):
        return self.avgpool(x)


# 定义网络配置(cfg)来构建完整的网络
cfg = [
    ('ConvBNReLU', 3, 64, 3, 1),  # 传递4个参数
    ('MaxPooling', 2, 2, 0),
    ('ConvBNReLU', 64, 128, 3, 1),
    ('MaxPooling', 2, 2, 0),
    ('ConvBNReLU', 128, 256, 3, 1),
    ('AvgPooling', 4, 1, 0),
]


# 构建网络
class CustomNet(nn.Module):
    def __init__(self, cfg):
        super(CustomNet, self).__init__()
        self.layers = nn.ModuleList()
        in_channels = 3  # 输入通道数

        for layer_cfg in cfg:
            layer_name, *layer_params = layer_cfg
            layer = layer_register.get_layer(layer_name)

            if layer_name in ['ConvBNReLU', 'BatchNorm2d']:
                self.layers.append(layer(in_channels, *layer_params))
                in_channels = layer_params[1]
            else:
                self.layers.append(layer(*layer_params))



# 创建完整的网络实例
custom_net = CustomNet(cfg)

# 打印网络结构
print(custom_net)
